import cv2
import numpy as np
import matplotlib.pyplot as plt

def CV_show(name, img):
    """显示图像并保持窗口，按任意键关闭"""
    cv2.imshow(name, img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

def sort_contours(cnts, method="left-to-right"):
    """
    对轮廓进行排序（支持 left-to-right / bottom-to-top / right-to-left）
    :param cnts: 轮廓列表
    :param method: 排序方式
    :return: 排序后的轮廓、边界框
    """
    reverse = False
    i = 0
    if method == "bottom-to-top":
        reverse = True
        i = 1
    elif method == "right-to-left":
        reverse = True
        i = 0

    boundingBoxes = [cv2.boundingRect(c) for c in cnts]
    (cnts, boundingBoxes) = zip(*sorted(zip(cnts, boundingBoxes),
                                        key=lambda b: b[1][i], reverse=reverse))
    return cnts, boundingBoxes

def resize(image, width=300):
    """
    按宽度缩放图像（保持宽高比）
    :param image: 输入图像
    :param width: 目标宽度
    :return: 缩放后的图像
    """
    (h, w) = image.shape[:2]
    r = width / float(w)
    new_h = int(h * r)
    return cv2.resize(image, (width, new_h), interpolation=cv2.INTER_AREA)


# -------------------------- 1. 模板图像预处理（数字模板） --------------------------
# 读取模板图像（灰度模式）
template_image = cv2.imread('152_1753154394_hd.jpeg', cv2.IMREAD_GRAYSCALE)
# 二值化（Otsu 自动阈值，反相）
ret, ref = cv2.threshold(template_image, 10, 255, cv2.THRESH_BINARY_INV)
# 提取轮廓（新版 OpenCV 返回 2 个值）
contours, hierarchy = cv2.findContours(ref.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
print(f"模板图像提取到 {len(contours)} 个轮廓")

# 存储数字模板（键: 索引，值: 归一化后的 ROI）
digits = {}
# 对轮廓排序后再处理
sorted_contours, _ = sort_contours(contours, "left-to-right")
for (i, c) in enumerate(sorted_contours):
    x, y, w, h = cv2.boundingRect(c)
    roi = template_image[y:y + h, x:x + w]
    # 统一调整为 (57, 58) 尺寸
    roi = cv2.resize(roi, (57, 58))
    digits[i] = roi


# -------------------------- 2. 目标图像预处理（待识别图像） --------------------------
# 读取目标图像（灰度模式）
target_image = cv2.imread('149_1753154392_hd.jpeg', cv2.IMREAD_GRAYSCALE)
# 缩放图像（宽度固定为 300，保持宽高比）
target_image = resize(target_image, 300)

# 形态学 TopHat 运算（突出亮区域）
rect_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 3))
tophat = cv2.morphologyEx(target_image, cv2.MORPH_TOPHAT, rect_kernel)

# Sobel 水平梯度计算
grad_x = cv2.Sobel(tophat, cv2.CV_32F, dx=1, dy=0, ksize=-1)
grad_x = np.absolute(grad_x)
(min_val, max_val) = (np.min(grad_x), np.max(grad_x))
# 归一化到 0-255
grad_x = (255 * ((grad_x - min_val) / (max_val - min_val))).astype("uint8")

# 形态学闭运算（连接水平区域）
sq_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (10, 5))
grad_x_closed = cv2.morphologyEx(grad_x, cv2.MORPH_CLOSE, sq_kernel)

# 二值化（Otsu 自动阈值）
ret, thresh = cv2.threshold(grad_x_closed, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
# 多次闭运算强化（可根据实际效果调整次数）
for _ in range(3):
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, sq_kernel)

# 提取轮廓
cnts, _ = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)


# -------------------------- 3. 筛选轮廓并识别数字 --------------------------
cur_image = target_image.copy()
# 转换为彩色图用于绘制（方便显示红色框、文字）
cur_image_color = cv2.cvtColor(cur_image, cv2.COLOR_GRAY2BGR)

locs = []
# 筛选轮廓（宽高比、尺寸过滤）
for (i, c) in enumerate(cnts):
    (x, y, w, h) = cv2.boundingRect(c)
    ar = w / float(h)
    # 根据宽高比和尺寸筛选
    if 2.5 < ar < 4.0 and 40 < w < 55 and 10 < h < 20:
        locs.append((x, y, w, h))

# 按 x 坐标排序（从左到右）
locs = sorted(locs, key=lambda x: x[0])
output = []

for (i, (gx, gy, gw, gh)) in enumerate(locs):
    group_output = []
    # 安全截取区域（避免越界）
    top = max(0, gy - 5)
    bottom = min(cur_image.shape[0], gy + gh + 5)
    left = max(0, gx - 5)
    right = min(cur_image.shape[1], gx + gw + 5)
    group = cur_image[top:bottom, left:right]

    # 二值化（Otsu 自动阈值）
    ret, group = cv2.threshold(group, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    # 提取轮廓
    bianyuan, _ = cv2.findContours(group.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # 排序轮廓（从左到右）
    bianyuan, _ = sort_contours(bianyuan, "left-to-right")

    for (j, c) in enumerate(bianyuan):
        if len(c) == 0:
            continue  # 跳过空轮廓
        # 转换为 OpenCV 兼容格式
        c = np.array(c, dtype=np.int32)
        x, y, w, h = cv2.boundingRect(c)

        # 安全截取 ROI（基于模板图像）
        if y + h > template_image.shape[0] or x + w > template_image.shape[1]:
            continue
        roi = template_image[y:y + h, x:x + w]
        # 统一调整尺寸
        roi = cv2.resize(roi, (57, 58))

        scores = []
        # 模板匹配
        for (digit, digit_roi) in digits.items():
            result = cv2.matchTemplate(roi, digit_roi, cv2.TM_CCOEFF)
            (_, score, _, _) = cv2.minMaxLoc(result)
            scores.append(score)

        if scores:
            # 取最匹配的数字
            group_output.append(str(np.argmax(scores)))

    # 绘制边界框和识别结果
    if group_output:
        cv2.rectangle(cur_image_color, (left, top), (right, bottom), (0, 0, 255), 2)
        cv2.putText(cur_image_color, "".join(group_output), (gx, gy - 15),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 0, 255), 2)
        output.extend(group_output)


# -------------------------- 4. 结果输出 --------------------------
print("识别结果：", "".join(output))
CV_show('识别结果', cur_image_color)